package io.modifier.basic.emo.scheme.ksp

import io.modifier.basic.emo.scheme.AbstractSchemeDefStorage
import io.modifier.basic.emo.scheme.ActivityScheme
import io.modifier.basic.emo.scheme.ComposeScheme
import io.modifier.basic.emo.scheme.SchemeBoolArg
import io.modifier.basic.emo.scheme.SchemeBoolArgParser
import io.modifier.basic.emo.scheme.SchemeDef
import io.modifier.basic.emo.scheme.SchemeFloatArg
import io.modifier.basic.emo.scheme.SchemeFloatArgParser
import io.modifier.basic.emo.scheme.SchemeIntArg
import io.modifier.basic.emo.scheme.SchemeIntArgParser
import io.modifier.basic.emo.scheme.SchemeLongArg
import io.modifier.basic.emo.scheme.SchemeLongArgParser
import io.modifier.basic.emo.scheme.SchemeStringArg
import io.modifier.basic.emo.scheme.SchemeStringArgParser
import com.google.devtools.ksp.KspExperimental
import com.google.devtools.ksp.getAnnotationsByType
import com.google.devtools.ksp.processing.CodeGenerator
import com.google.devtools.ksp.processing.Dependencies
import com.google.devtools.ksp.processing.KSPLogger
import com.google.devtools.ksp.processing.Resolver
import com.google.devtools.ksp.processing.SymbolProcessor
import com.google.devtools.ksp.symbol.KSAnnotated
import com.google.devtools.ksp.symbol.KSClassDeclaration
import com.google.devtools.ksp.symbol.KSDeclaration
import com.google.devtools.ksp.symbol.KSFunctionDeclaration
import com.google.devtools.ksp.symbol.KSType
import com.google.devtools.ksp.symbol.Nullability
import java.io.OutputStream

class SchemeProcessor(
    private val codeGenerator: CodeGenerator,
    private val logger: KSPLogger,
) : SymbolProcessor {

    override fun process(resolver: Resolver): List<KSAnnotated> {
        val activityList = resolver
            .getSymbolsWithAnnotation(ActivityScheme::class.java.name)
            .filterIsInstance<KSClassDeclaration>()
            .toSet()
            .sortedBy { it.simpleName.getShortName() }

        val composeList = resolver
            .getSymbolsWithAnnotation(ComposeScheme::class.java.name)
            .filterIsInstance<KSFunctionDeclaration>()
            .toSet()
            .sortedBy { it.simpleName.getShortName() }

        if (activityList.isEmpty() && composeList.isEmpty()) {
            return emptyList()
        }
        val composeGraph = processSchemeStorage(activityList, composeList)
        processComposeGraphBuilder(composeGraph)

        return emptyList()
    }

    @OptIn(KspExperimental::class)
    private fun processSchemeStorage(
        activityList: List<KSClassDeclaration>,
        composeList: List<KSFunctionDeclaration>,
    ): List<ComposeGraphItem> {
        val storageInputFlies = sequenceOf(activityList, composeList)
            .flatMap { it.asSequence() }
            .map { it.containingFile!! }
            .toSet()
            .toTypedArray()
        val storagePackageName = "io.modifier.basic.emo.scheme.impl"
        val osForStorage: OutputStream = codeGenerator.createNewFile(
            dependencies = Dependencies(true, *storageInputFlies),
            packageName = storagePackageName,
            fileName = "GeneratedSchemeDefStorage"
        )
        val composeGraph = mutableListOf<ComposeGraphItem>()
        osForStorage.writeLine("package $storagePackageName")
        osForStorage.writeLine("import androidx.annotation.Keep")
        osForStorage.writeLine("import io.modifier.basic.emo.scheme.AbstractSchemeDefStorage")
        osForStorage.writeLine("import io.modifier.basic.emo.scheme.SchemeDef")
        osForStorage.writeLine("import io.modifier.basic.emo.scheme.SchemeArgDefine")
        osForStorage.writeLine("@Keep")
        osForStorage.write("class GeneratedSchemeDefStorage:${AbstractSchemeDefStorage::class.java.simpleName}()")
        osForStorage.writeBlock {
            osForStorage.write("init")
            osForStorage.writeBlock {
                var nextSchemeId = 1
                activityList.asSequence().forEach { cls ->
                    val schemes = cls.getAnnotationsByType(ActivityScheme::class)
                    val args = cls.buildArgDefineList()
                    schemes.forEach {
                        writeLine(
                            "add(SchemeDef(${nextSchemeId++}, \"${it.action}\", " +
                                    "emptyList(), $args, \"${cls.qualifiedName!!.asString()}\"," +
                                    "${it.transition}))"
                        )
                    }
                }

                composeList.asSequence().forEach { fn ->
                    // TODO can not use getAnnotationsByType because of usage of kClass. Just wait ksp's new version and check.
                    val schemes = fn.annotations.filter {
                        it.shortName.getShortName() == ComposeScheme::class.simpleName && it.annotationType.resolve().declaration
                            .qualifiedName?.asString() == ComposeScheme::class.qualifiedName
                    }
                    val args = fn.buildArgDefineList()
                    schemes.forEach {
                        val action = it.arguments.find { arg -> arg.name!!.asString() == "action" }!!.value as String
                        val host = it.arguments.find { arg -> arg.name!!.asString() == "alternativeHosts" }!!.value as List<*>
                        val transition = it.arguments.find { arg -> arg.name!!.asString() == "transition" }!!.value as Int
                        if (host.isEmpty()) {
                            throw RuntimeException("ComposeScheme.alternativeHosts for ${fn.simpleName.getShortName()} can not be empty")
                        }
                        val alternativeHosts = host.joinToString(",") { h ->
                            "${(h as KSType).declaration.qualifiedName!!.asString()}::class"
                        }
                        writeLine(
                            "add(SchemeDef($nextSchemeId, \"${action}\", " +
                                    "listOf($alternativeHosts), $args, \"${SchemeDef.COMPOSE_CLASS_SUFFIX}\", " +
                                    "$transition))"
                        )
                        host.forEach { h ->
                            composeGraph.add(ComposeGraphItem(fn, h as KSType, nextSchemeId))
                        }
                        nextSchemeId++
                    }
                }
            }
        }
        osForStorage.close()
        return composeGraph
    }

    @OptIn(KspExperimental::class)
    private fun KSDeclaration.buildArgDefineList(): String {
        val boolArg = getAnnotationsByType(SchemeBoolArg::class)
            .map {
                "SchemeArgDefine(\"${it.name}\", ${it.special}, ${SchemeBoolArgParser::class.qualifiedName}, ${it.default})"
            }
        val intArg = getAnnotationsByType(SchemeIntArg::class)
            .map {
                "SchemeArgDefine(\"${it.name}\", ${it.special}, ${SchemeIntArgParser::class.qualifiedName}, ${it.default})"
            }
        val longArg = getAnnotationsByType(SchemeLongArg::class)
            .map {
                "SchemeArgDefine(\"${it.name}\", ${it.special}, ${SchemeLongArgParser::class.qualifiedName}, ${it.default})"
            }
        val floatArg = getAnnotationsByType(SchemeFloatArg::class)
            .map {
                "SchemeArgDefine(\"${it.name}\", ${it.special}, ${SchemeFloatArgParser::class.qualifiedName}, ${it.default}f)"
            }
        val stringArg = getAnnotationsByType(SchemeStringArg::class)
            .map {
                "SchemeArgDefine(\"${it.name}\", ${it.special}, ${SchemeStringArgParser::class.qualifiedName}, \"${it.default}\")"
            }
        val list = sequenceOf(boolArg, intArg, longArg, floatArg, stringArg)
            .flatMap { it.asSequence() }
            .joinToString(",")
        return "listOf($list)"
    }

    private fun processComposeGraphBuilder(
        composeGraph: List<ComposeGraphItem>,
    ) {
        composeGraph.groupBy {
            it.host
        }.forEach { (host, items) ->
            val inputFiles = items.asSequence()
                .map { it.fn.containingFile!! }
                .toSet()
                .toTypedArray()
            val packageName = host.declaration.qualifiedName!!.asString().let { it.substring(0, it.lastIndexOf('.')) }
            val clsName = host.declaration.simpleName.asString() + SchemeDef.COMPOSE_CLASS_SUFFIX
            val os: OutputStream = codeGenerator.createNewFile(
                dependencies = Dependencies(true, *inputFiles),
                packageName = packageName,
                fileName = clsName
            )

            os.writeLine("package $packageName")
            os.writeLine("import androidx.annotation.Keep")
            os.writeLine("import androidx.navigation.NavGraphBuilder")
            os.writeLine("import androidx.compose.runtime.Composable")
            os.writeLine("import androidx.compose.animation.ExperimentalAnimationApi")
            os.writeLine("import androidx.navigation.compose.composable")
            os.writeLine("import io.modifier.basic.emo.scheme.SchemeClient")
            os.writeLine("import io.modifier.basic.emo.scheme.impl.toComposeRouteDefine")
            os.writeLine("import io.modifier.basic.emo.scheme.impl.toComposeNavArg")
            os.writeLine("import io.modifier.basic.emo.scheme.impl.ComposeSchemeNavGraphBuilder")
            os.writeLine("import io.modifier.basic.emo.scheme.impl.SchemeTransitionProviders")
            os.writeLine("import io.modifier.basic.emo.scheme.parseModelData")
            os.writeLine("@Keep")
            os.write("class $clsName: ComposeSchemeNavGraphBuilder")
            os.writeBlock {
                os.writeLine("@OptIn(ExperimentalAnimationApi::class)")
                os.write(
                    "override fun build(" +
                            "client: SchemeClient, " +
                            "navGraphBuilder: NavGraphBuilder)"
                )
                os.writeBlock {
                    items.forEach { item ->
                        os.write("client.storage.findById(${item.schemeId})!!.let")
                        os.writeBlock {
                            os.writeLine("val transition = SchemeTransitionProviders.get(it.transition)")
                            os.write(
                                "navGraphBuilder.composable(" +
                                        "it.toComposeRouteDefine()," +
                                        "it.toComposeNavArg()," +
                                        "enterTransition = transition.enterTransition()," +
                                        "exitTransition = transition.exitTransition()," +
                                        "popEnterTransition = transition.popEnterTransition()," +
                                        "popExitTransition = transition.popExitTransition()" +
                                        "){"
                            )
                            if (item.fn.parameters.size == 1) {
                                os.writeLine(" entry ->")
                                val param = item.fn.parameters[0].type.resolve()
                                val name = param.declaration.qualifiedName!!.asString()
                                if (name == "androidx.navigation.NavBackStackEntry") {
                                    os.writeLine("${item.fn.qualifiedName!!.asString()}(entry)")
                                } else {
                                    os.writeLine("val model = entry.parseModelData<$name>()")
                                    if (param.nullability == Nullability.NOT_NULL) {
                                        os.writeLine("if(model != null){${item.fn.qualifiedName!!.asString()}(model)}")
                                    } else {
                                        os.writeLine("${item.fn.qualifiedName!!.asString()}(model)")
                                    }
                                }
                            } else if (item.fn.parameters.size == 2) {
                                os.writeLine(" entry ->")
                                val param = item.fn.parameters[1].type.resolve()
                                val name = param.declaration.qualifiedName!!.asString()
                                os.writeLine("val model = entry.parseModelData<$name>()")
                                if (param.nullability == Nullability.NOT_NULL) {
                                    os.writeLine("if(model!=null){${item.fn.qualifiedName!!.asString()}(entry, model)}")
                                } else {
                                    os.writeLine("${item.fn.qualifiedName!!.asString()}(entry, model)")
                                }
                            } else if (item.fn.parameters.isEmpty()) {
                                os.writeLine("${item.fn.qualifiedName!!.asString()}()")
                            } else {
                                throw RuntimeException(
                                    "${item.fn.simpleName} can have one param with type NavBackStackEntry or have any param."
                                )
                            }
                            os.writeLine("}")
                        }
                    }
                }
            }
            os.close()
        }
    }

    private fun OutputStream.write(str: String) = write(str.toByteArray())
    private fun OutputStream.writeLine(str: String) = write(str + "\n")

    private fun OutputStream.writeBlock(block: OutputStream.() -> Unit) {
        writeLine("{")
        block()
        writeLine("}")
    }
}

class ComposeGraphItem(
    val fn: KSFunctionDeclaration,
    val host: KSType,
    val schemeId: Int,
)
